from langchain.agents import AgentExecutor,create_tool_calling_agent
from tool.neo4j_tool import Neo4jTool
from tool.optimized_chroma_tool import ChromaTool
from tool.get_datetime_now import GetTimeTool
from model.my_chat_model import ChatModel
from langchain_core.prompts import ChatPromptTemplate
from langchain.memory import ConversationBufferWindowMemory
from optimizer.query_rewriter import rewrite_query

# 全局变量：存储对话记忆（适用于单用户）
memory = ConversationBufferWindowMemory(k=5, return_messages=True)
# 工具只需创建一次
tools = [Neo4jTool(),ChromaTool(),GetTimeTool()]
#创建一个聊天的智能体
def create_agent(question):
    #查询重写：语义补全 + 意图澄清
    rewritten_question = rewrite_query(question)
    print(f"原始问题：{question}")
    print(f"重写问题:{rewritten_question}")
    #获取大模型
    chat=ChatModel()
    llm=chat.get_line_model()
    #创建提示词模板
    prompt=ChatPromptTemplate([
        ("system", """
             你是一个智能的AI查询助手，你可以使用以下几个工具
                 1 get_neo4j_tool 用于查询股东、高管、公司之间的持股、投资、关联、任职、受益及风险等问
                 2 get_chroma_tool 主要用于查询chroma数据库的企业背景数据
                 3 get_time_tool 主要用于查询当前时刻的时间
                 如果用户的提问的问题是关于股东、高管、公司之间的持股、投资、关联、任职、受益以及风险等问题时，请执行工具 get_neo4j_tool；
                 如果用户提问的问题是有关企业背景的调查问题，请执行工具 get_chroma_tool；请给用户准确信息，其他文本不需要
                 如果用户提问的问题是有关当前时刻时间的获取，请执行工具get_time_tool
            """),
        ("placeholder", "{history}"),  # 必须有
        ("human","{input}"),
        ("placeholder","{agent_scratchpad}") #agent的思考过程
    ])
    #创建工具调用agent
    agent=create_tool_calling_agent(llm,tools,prompt)
    #创建智能体执行器
    agent_executor=AgentExecutor(
        agent=agent,
        tools=tools,
        memory=memory,
        verbose=True
    )
    rs=agent_executor.invoke({"input":rewritten_question})
    print(rs["output"])
    return rs["output"]

if __name__ == '__main__':
    create_agent("公司名称为公司15的公司怎么样？")
